import os
import sys

# 添加项目根目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(os.path.dirname(current_dir))
sys.path.insert(0, project_root)

from devdeploy.inference.inference import create_inference
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt

# 模型路径和图片目录
CHECKPOINT_PATH = '/Users/zhanghaining/git/SteelAiLab/out_models/unet_deeplabv3_no_crop_dark/fullmodel_iter_best_mDice_99_7400.onnx'
# CHECKPOINT_PATH = 'out_models/swin_tiny/fullmodel_final.pth'
IMG_DIR = '/Volumes/data1/JH/dataset/datadark/all_data/train/img'

# 创建推理器实例（分割任务不需要class_names）
# infer = Inference(CHECKPOINT_PATH, device='cpu')
infer = create_inference(CHECKPOINT_PATH, device='cpu')

def test_pil_image_input():
    """测试使用PIL Image对象作为输入的分割推理功能"""
    print("\n" + "="*80)
    print("PIL Image输入分割推理测试")
    print("="*80)
    
    # 获取所有测试图片
    image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif')
    test_images = [
        f for f in os.listdir(IMG_DIR) 
        if f.lower().endswith(image_extensions)
    ]
    
    if not test_images:
        print(f"错误：在目录 {IMG_DIR} 中没有找到测试图片")
        return
    
    print(f"找到 {len(test_images)} 张测试图片")
    print("-" * 80)
    
    # 对每张图片进行测试
    for i, test_image in enumerate(test_images, 1):
        img_path = os.path.join(IMG_DIR, test_image)
        
        try:
            # 打开图片
            img = Image.open(img_path)
            
            # 使用PIL Image对象进行分割推理
            result = infer.infer_single(img)
            
            # 打印结果
            print(f'图片: {test_image}')
            print(f'  分割掩码形状: {result["segmentation_mask"].shape}')
            print(f'  类别数量: {result["num_classes"]}')
            print(f'  掩码值范围: [{result["segmentation_mask"].min()}, {result["segmentation_mask"].max()}]')
            
            # 显示分割结果
            # if i <= 3:  # 只显示前3张图片的结果
            display_segmentation_result(img, result, test_image)
            
            print("-" * 40)
            
        except Exception as e:
            print(f"测试失败 {test_image}: {str(e)}")

def display_segmentation_result(original_img, result, img_name):
    """显示分割结果"""
    try:
        # 创建图形
        fig, axes = plt.subplots(1, 2, figsize=(12, 6))
        
        # 显示原图
        axes[0].imshow(original_img)
        axes[0].set_title(f'Original Image: {img_name}')
        axes[0].axis('off')
        
        # 显示分割掩码
        seg_mask = result['segmentation_mask']
        axes[1].imshow(seg_mask, cmap='gray')
        axes[1].set_title(f'Segmentation Mask (Shape: {seg_mask.shape})')
        axes[1].axis('off')
        
        plt.tight_layout()
        plt.show()
        
    except Exception as e:
        print(f"显示分割结果失败: {str(e)}")

def test_model_info():
    """测试模型信息获取功能"""
    print(f"\n模型信息:")
    print("-" * 40)
    
    try:
        # 获取任务类型
        task_type = infer.get_task_type()
        print(f"任务类型: {task_type}")
        
        # 获取模型信息
        model_info = infer.get_model_info()
        print(f"模型信息: {model_info}")
        
        # 验证模型
        is_valid = infer.validate_model()
        print(f"模型验证: {'通过' if is_valid else '失败'}")
        
        # 显示图像尺寸信息
        if hasattr(infer, 'img_scale') and infer.img_scale:
            print(f"图像尺寸: {infer.img_scale}")
        
    except Exception as e:
        print(f"获取模型信息失败: {str(e)}")

def test_single_image_inference():
    """测试单张图片分割推理功能"""
    print("\n" + "="*80)
    print("单张图片分割推理测试")
    print("="*80)
    
    # 获取测试图片列表
    image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif')
    test_images = [
        f for f in os.listdir(IMG_DIR) 
        if f.lower().endswith(image_extensions)
    ]
    
    if not test_images:
        print(f"错误：在目录 {IMG_DIR} 中没有找到测试图片")
        return
    
    print(f"找到 {len(test_images)} 张测试图片")
    print("-" * 60)
    
    # 测试所有图片的单张推理
    for i, img_name in enumerate(test_images):
        img_path = os.path.join(IMG_DIR, img_name)
        
        try:
            # 单张图片推理
            result = infer.infer_single(img_path)
            
            # 打印结果
            print(f'图片: {img_name}')
            print(f'  分割掩码形状: {result["segmentation_mask"].shape}')
            print(f'  类别数量: {result["num_classes"]}')
            print(f'  掩码值范围: [{result["segmentation_mask"].min()}, {result["segmentation_mask"].max()}]')
            
        except Exception as e:
            print(f"推理失败 {img_name}: {str(e)}")

def test_batch_inference():
    """测试批量分割推理功能"""
    print("\n" + "="*80)
    print("批量分割推理测试")
    print("="*80)
    
    try:
        # 批量推理
        results = infer.infer_batch(IMG_DIR)

        # 打印结果
        for img_name, result in results.items():
            if 'error' in result:
                print(f'错误: {img_name} - {result["error"]}')
            else:
                print(f'图片: {img_name}')
                print(f'  分割掩码形状: {result["segmentation_mask"].shape}')
                print(f'  类别数量: {result["num_classes"]}')
                print(f'  掩码值范围: [{result["segmentation_mask"].min()}, {result["segmentation_mask"].max()}]')
                print("-" * 30)
                
    except Exception as e:
        print(f"批量推理失败: {str(e)}")

# 执行测试
if __name__ == "__main__":
    # 测试模型信息
    test_model_info()
    
    # 测试PIL Image输入
    test_pil_image_input()
    
    # 测试单张图片推理
    # test_single_image_inference()
    
    # 测试批量推理
    # test_batch_inference()
    
    print("\n" + "="*80)
    print("分割推理测试完成")
    print("="*80) 